#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.0
# @author : cd
# @Email : 19688513@qq.com
# @Project : new-horizons-engine
# @File : StockPlotter.py
# @Software: PyCharm
# @time: 2025/6/4 10:19
# @description : 股票图表绘制器类，负责绘制股票分析图表
"""

import urllib3
import logging
from typing import List, Tuple, Dict, Any, Optional
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.dates import DateFormatter, DayLocator
import os

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# 禁用不安全的请求警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


class StockPlotter:
    """股票图表绘制器类，负责绘制股票分析图表"""

    def __init__(self):
        # 设置中文字体支持
        plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'WenQuanYi Micro Hei']
        plt.rcParams['axes.unicode_minus'] = False

        # 检查系统并设置字体
        if os.name == 'nt':
            # Windows系统使用Microsoft YaHei字体
            plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
        else:
            # Linux/Mac系统使用WenQuanYi Zen Hei
            plt.rcParams['font.sans-serif'] = ['WenQuanYi Zen Hei']

    def plot_stock_analysis(self, df: pd.DataFrame, symbol: str, start_date: str, end_date: str) -> Optional[str]:
        """绘制股票分析图表"""
        if df.empty:
            logging.warning("无法绘制图表，数据为空")
            return None

        try:
            # 创建更大的图表（4行1列）
            plt.figure(figsize=(16, 24))
            plt.suptitle(f"{symbol} 股票分析 ({start_date} 至 {end_date})", fontsize=16, y=0.98)

            # 1. 绘制价格图表
            ax1 = plt.subplot(4, 1, 1)
            self._plot_price_chart(ax1, df)

            # 2. 绘制成交量图表
            ax2 = plt.subplot(4, 1, 2, sharex=ax1)
            self._plot_volume_chart(ax2, df)

            # 3. 绘制成交量MACD指标
            ax3 = plt.subplot(4, 1, 3, sharex=ax1)
            self._plot_macd_chart(ax3, df)

            # 4. 绘制其他技术指标
            ax4 = plt.subplot(4, 1, 4, sharex=ax1)
            self._plot_other_indicators(ax4, df)

            # 优化日期显示（只在最底部图表显示）
            plt.gca().xaxis.set_major_locator(DayLocator(interval=7))  # 每周显示一个日期
            plt.gca().xaxis.set_major_formatter(DateFormatter("%Y-%m-%d"))
            plt.gcf().autofmt_xdate()  # 自动旋转日期标签

            # 自动调整布局
            plt.tight_layout(rect=[0, 0, 1, 0.96])  # 为总标题留出空间

            # 保存图片
            plot_filename = f"stock_analysis_{symbol}_{start_date}_{end_date}.png"
            plt.savefig(plot_filename, dpi=150, bbox_inches='tight')
            plt.close()

            return plot_filename

        except ImportError:
            logging.warning("未安装matplotlib，无法生成价格走势图")
            return None
        except Exception as e:
            logging.error(f"绘图时出错: {str(e)}")
            return None

    def _plot_price_chart(self, ax, df: pd.DataFrame):
        """绘制价格图表"""
        # 绘制收盘价
        ax.plot(df['date'], df['close'], 'b-', label='收盘价', linewidth=1.5)

        # 定义价格均线标签和颜色
        price_labels = {
            'ma5': '5日均线',
            'ma10': '10日均线',
            'ma20': '20日均线',
            'ma30': '30日均线',
            'ma60': '60日均线',
            'ma120': '120日均线'
        }

        price_colors = {
            'ma5': 'green',
            'ma10': 'blue',
            'ma20': 'red',
            'ma30': 'purple',
            'ma60': 'orange',
            'ma120': 'brown'
        }

        # 绘制所有存在的价格均线
        for ma in ['ma5', 'ma10', 'ma20', 'ma30', 'ma60', 'ma120']:
            if ma in df.columns:
                ax.plot(df['date'], df[ma],
                        '-',
                        color=price_colors[ma],
                        linewidth=1.2,
                        label=price_labels[ma])

        # 设置标题和图例
        ax.set_title("价格走势", fontsize=12)
        ax.legend(loc='best', fontsize=8)
        ax.grid(True, linestyle='--', alpha=0.7)
        ax.set_ylabel('价格')

    def _plot_volume_chart(self, ax, df: pd.DataFrame):
        """绘制成交量图表"""
        # 定义成交量相关元素
        volume_labels = {
            'volume': '成交量(万手)',
            'vol_ma5': '5日均量',
            'vol_ma10': '10日均量',
            'vol_ma20': '20日均量',
            'vol_ma30': '30日均量',
            'vol_ma60': '60日均量',
            'vol_ma120': '120日均量'
        }

        volume_styles = {
            'volume': '-',  # 原始成交量使用实线
            'vol_ma5': '--',
            'vol_ma10': '-.',
            'vol_ma20': ':',
            'vol_ma30': '-',
            'vol_ma60': '--',
            'vol_ma120': '-.'
        }

        volume_colors = {
            'volume': 'gray',  # 原始成交量使用灰色
            'vol_ma5': 'red',
            'vol_ma10': 'blue',
            'vol_ma20': 'green',
            'vol_ma30': 'purple',
            'vol_ma60': 'orange',
            'vol_ma120': 'brown'
        }

        # 绘制原始成交量线图
        if 'volume' in df.columns:
            ax.plot(df['date'], df['volume'] / 10000,
                    volume_styles['volume'],
                    color=volume_colors['volume'],
                    linewidth=1.0,
                    alpha=0.7,
                    label=volume_labels['volume'])

        # 绘制所有存在的均量线
        for ma in ['vol_ma5', 'vol_ma10', 'vol_ma20', 'vol_ma30', 'vol_ma60', 'vol_ma120']:
            if ma in df.columns:
                ax.plot(df['date'], df[ma],
                        volume_styles[ma],
                        color=volume_colors[ma],
                        linewidth=1.5,
                        label=volume_labels[ma])

        ax.legend(loc='upper left', fontsize=8)
        ax.grid(True, linestyle='--', alpha=0.5)
        ax.set_title('成交量及均量线', fontsize=12)
        ax.set_ylabel('成交量(万手)')

    def _plot_macd_chart(self, ax, df: pd.DataFrame):
        """绘制MACD指标图表"""
        if 'vol_macd' in df.columns and 'vol_signal' in df.columns:
            # 绘制MACD和信号线
            ax.plot(df['date'], df['vol_macd'], 'b-', label='成交量MACD', linewidth=1.2)
            ax.plot(df['date'], df['vol_signal'], 'r-', label='信号线', linewidth=1.2)

            # 绘制MACD柱状图
            if 'vol_hist' in df.columns:
                # 确定柱状图颜色（正值为绿，负值为红）
                colors = ['g' if h >= 0 else 'r' for h in df['vol_hist']]
                ax.bar(df['date'], df['vol_hist'], color=colors, alpha=0.5, label='柱状图')

            ax.axhline(y=0, color='gray', linestyle='-', alpha=0.3)
            ax.legend(loc='best', fontsize=8)
            ax.grid(True, linestyle='--', alpha=0.5)
            ax.set_title('成交量MACD指标', fontsize=12)
            ax.set_ylabel('MACD值')
        else:
            ax.set_title('成交量MACD指标 - 数据不可用', fontsize=12)

    def _plot_other_indicators(self, ax, df: pd.DataFrame):
        """绘制其他技术指标图表"""
        # 绘制成交量RSI
        if 'vol_rsi' in df.columns:
            ax.plot(df['date'], df['vol_rsi'], 'g-', label='成交量RSI', linewidth=1.5)
            ax.axhline(y=30, color='r', linestyle='--', alpha=0.5)
            ax.axhline(y=70, color='r', linestyle='--', alpha=0.5)

        # 绘制成交量布林带
        if 'vol_boll_upper' in df.columns:
            ax.plot(df['date'], df['vol_boll_upper'], 'b--', label='布林带上轨', alpha=0.7)
            ax.plot(df['date'], df['vol_boll_middle'], 'k-', label='布林带中轨', alpha=0.7)
            ax.plot(df['date'], df['vol_boll_lower'], 'b--', label='布林带下轨', alpha=0.7)
            ax.fill_between(df['date'],
                            df['vol_boll_upper'],
                            df['vol_boll_lower'],
                            color='gray', alpha=0.1)

        # 绘制成交量KDJ
        if 'vol_k' in df.columns and 'vol_d' in df.columns and 'vol_j' in df.columns:
            ax.plot(df['date'], df['vol_k'], 'm-', label='K线', linewidth=1.2)
            ax.plot(df['date'], df['vol_d'], 'c-', label='D线', linewidth=1.2)
            ax.plot(df['date'], df['vol_j'], 'y-', label='J线', linewidth=1.2)

        # 设置图例和标题
        ax.legend(loc='best', fontsize=8)
        ax.grid(True, linestyle='--', alpha=0.5)
        ax.set_title('成交量技术指标 (RSI, 布林带, KDJ)', fontsize=12)
        ax.set_ylabel('指标值')

        # 添加零线
        ax.axhline(y=0, color='gray', linestyle='-', alpha=0.3)